#!/usr/bin/env python
# -*- coding: utf-8 -*-

# @Time   : 2021/1/17 11.17
# @Author : Ray'
# @File   : Sniff_P2P.py

# hdbnoc password ：njnet@seu

from scapy.all import *
import datetime
import SSHConnection
import Win_Soft
import threading
import os


## ========================================================================== ##
## ------------------------ USER CONFIGURATION FLAGS ------------------------ ##
## ========================================================================== ##

catch_traffic_time_set = 60

## ==================  Local Machine Config  ================== ##
# ipV4
local_ip = 'XXX.XXX.XXX.XXX'

# .torrent files folder
Seeds = "D:\\Ray_Xiao_c\\torrent\\"

# Directory for saving pcap
dir_save_pcap = "D:\\Ray_Xiao_c\\pcap\\tor_p2p\\"

# Proxy
proxy = {'host': '127.0.0.1', 'port': 9050}

# Local Machine Network Interface
local_iface = 'Intel(R) Dual Band Wireless-AC XXXX'


## ==================  Server Config  ================== ##
## Gateway
gateway_dict = {'host': 'XXX.XXX.XXX.XXX', 'port': 2222, 'username': '', 'pwd': ''}
gateway_path = '/media/work/'

## Server ipv4
server_IP = 'XXX.XXX.XXX.XXX'



def catch_traffic(catch_traffic_time, num):
    """ Sniffing traffic in local machine
    :param catch_traffic_time
    :return: none
    """
    IPfilter = 'host ' + server_IP + ' and tcp'
    PTKS = sniff(filter=IPfilter, iface=local_iface, timeout=catch_traffic_time)
    pcapname = dir_save_pcap + str(num) + '\\' + datetime.datetime.now().strftime('%Y%m%d_%H%M_%S_') + '.pcap'
    wrpcap(pcapname, PTKS)


def catch_traffic_gateway(gatewaySSH, catch_traffic_time, num):
    """ Sniffing traffic in gateway
    :param gatewaySSH:  A object instantiated for connecting to remote gateway
    :param catch_traffic_time
    :return: none
    """
    gateway_cmd = 'sudo timeout ' + str(catch_traffic_time) + ' tcpdump host ' + \
                  server_IP + ' and ' + local_ip + ' -s0 -G ' +\
                  str(catch_traffic_time) + ' -w ' + gateway_path + '/' + str(num) + '/%Y_%m%d_%H%M_%S.pcap'
    gatewaySSH.cmd(gateway_cmd, sudo=True)



def capture_traffic(catch_traffic_time):
    """ main function
    :param catch_traffic_time:  the time for catching traffic
    :return: none
    """

    # Connecting to gateway
    gatewaySSH = SSHConnection.SSHConnection(gateway_dict)
    gatewaySSH.connect()

    for i in range(10):        # 143 .torrent files，loop 10 times

        os.mkdir(dir_save_pcap + str(i))

        g_path = gateway_path + '/' + str(i)
        gatewaySSH.cmd('sudo mkdir ' + g_path, sudo=True)

        for fold in os.listdir(Seeds):
            #print('fold编号 %s !!!' % fold)
            for torrent in os.listdir(os.path.join(Seeds, fold)):         # only one
                # ========================= Start manipulating Bittorrent =========================#
                bittorrent = Win_Soft.Win_Soft(os.path.join(Seeds, fold, torrent))
                bittorrent.manipulate_bittorrent()  # The OK button is not pressed at this time

                print("Start Sniffing")

                # ========================= Local machine start sniffing  =========================#
                catch_traffic_thread = threading.Thread(target=catch_traffic, args=(catch_traffic_time, i+11))
                catch_traffic_thread.start()

                # ========================= gateway start sniffing =========================#
                catch_gateway_traffic_thread = threading.Thread(target=catch_traffic_gateway,
                                                                args=(gatewaySSH, catch_traffic_time, i+11))
                catch_gateway_traffic_thread.start()

                print("Start downloading")

                ## pressing OK button here
                bittorrent.movePos(880, 620)
                time.sleep(1)
                bittorrent.click_left()

                catch_traffic_thread.join()
                catch_gateway_traffic_thread.join()  # Block for wating

                bittorrent.delete_task()  # Completely delete tasks from the list

                time.sleep(10)  # sleep 10s

    # close gateway ssh
    gatewaySSH.close()






if __name__ == "__main__":

    capture_traffic(catch_traffic_time_set)

    print("^-^success^-^")